package online.heycm.excel.excel;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.metadata.WriteSheet;
import org.apache.ibatis.session.ResultContext;
import org.apache.ibatis.session.ResultHandler;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;

/**
 * @author hey
 * @version 1.0
 * @since 2024/8/17 22:37
 */
public abstract class ExcelHandler<T, U> implements ResultHandler<T> {

    // 一个sheet最多写入 10w 条数据
    private final int batchSize;
    // sheet结果集
    private final List<U> result;
    // 行数
    private int row;
    // sheet
    private int sheet;
    // ExcelWriter
    private final ExcelWriter excelWriter;

    public ExcelHandler(String fileName) {
        this.batchSize = 10 * 10000;
        result = new ArrayList<>(batchSize);
        excelWriter = EasyExcel.write(fileName).build();
    }

    public ExcelHandler(int batchSize, String fileName) {
        this.batchSize = Math.min(batchSize, 10 * 10000);
        result = new ArrayList<>(batchSize);
        excelWriter = EasyExcel.write(fileName).build();
    }

    @Override
    public void handleResult(ResultContext<? extends T> resultContext) {
        T resultObject = resultContext.getResultObject();
        result.add(handleResult(resultObject));
        row++;
        if (row % batchSize == 0 || resultContext.isStopped()) {
            writeExcel();
        }
    }

    private void writeExcel() {
        if (result.isEmpty()) {
            return;
        }
        U u = result.get(0);
        try {
            WriteSheet writeSheet = EasyExcel.writerSheet(sheet++).head(u.getClass()).build();
            excelWriter.write(result, writeSheet);
        } catch (Exception e) {
            e.printStackTrace();
        }
        // for gc
        result.clear();
    }

    private void flush() {
        if (!result.isEmpty()) {
            writeExcel();
        }
    }

    public abstract U handleResult(T t);

    public void handle(Consumer<ExcelHandler<T, U>> consumer) {
        consumer.accept(this);
        flush();
        excelWriter.close();
    }

    public int getRow() {
        return row;
    }
}
